#!/usr/bin/python3.9
# code=utf-8

import tracemalloc


class MemoryTrace(object):
    def __init__(self, n_snapshots=25):
        self.n_snapshots = n_snapshots
        self.snapshots = {}
        tracemalloc.start(n_snapshots)

    def take_snapshot(self, name):

        snapshot = tracemalloc.take_snapshot()
        self.snapshots[name] = snapshot
        print(f" snapshot '{name}' saved")
        return snapshot

    def compare_snapshots(self, name1, name2):
        """compare two snapshot"""
        if name1 not in self.snapshots or name2 not in self.snapshots:
            print(f"error: {name1} or {name2} not exist")
            return

        snapshot1 = self.snapshots[name1]
        snapshot2 = self.snapshots[name2]

        stats = snapshot2.compare_to(snapshot1, 'lineno')

        print(f"\n compare '{name1}' -> '{name2}':")

        increased_stats = [stat for stat in stats if stat.size > 0 and 'tracemalloc.py' not in str(stat.traceback[0])  ]
        # and 'tracemalloc.py' not in str(stat.traceback[0])

        if increased_stats:
            print(f"position:")
            for i, stat in enumerate(increased_stats[:10], 1):
                # print(str(stat.traceback[0]))
                print(f"#{i}: +{stat.size / 1024:.1f} KiB - {stat.traceback[0]}")
        else:
            print(f"no increased memory")

        return increased_stats

    def detect_leaks(self, baseline_name, leak_name, threshold_kb=100):
        increased_stats = self.compare_snapshots(baseline_name, leak_name)

        if not increased_stats:
            print(" no memory leak")
            return False

        total_increase = sum(stat.size for stat in increased_stats) / 1024  # KB
        max_increase = max(stat.size for stat in increased_stats) / 1024  # KB

        print(f"\nresult:")
        print(f"  total: {total_increase:.1f} KB")
        print(f"  max item: {max_increase:.1f} KB")

        if total_increase > threshold_kb:
            print(f" memory leak (> {threshold_kb} KB)")
            return True
        else:
            print(" memory normal (< {threshold_kb} KB)")
            return False

    def clean_up(self):
        tracemalloc.stop()
        self.snapshots.clear()


def example():
    """内存泄漏检测示例"""
    # 1
    detector = MemoryTrace()

    detector.take_snapshot("baseline")


    # user code --------------------------------------
    # 2 normal operation
    normal_data = []
    for i in range(1000):
        normal_data.append(f"temp_{i}")

    detector.take_snapshot("normal_operation")

    # 2 leak operation
    leaked_data = []
    for i in range(10000):
        leaked_data.append(f"leak_{i}" * 10)

    detector.take_snapshot("leak_operation")

    # user code --------------------------------------

    # 3 check
    print("=== check ===")
    detector.detect_leaks("baseline", "normal_operation")
    detector.detect_leaks("baseline", "leak_operation")

    # 4 clean
    detector.clean_up()


if __name__ == '__main__':
    example()